#include <string.h>
#include <stdint.h>
#include <stdlib.h>
#include <stdio.h>

#include <psp2/kernel/processmgr.h>
#include <psp2/kernel/threadmgr.h>
#include <psp2/kernel/sysmem.h>
#include <psp2/audioout.h>
#include <speex/speex_resampler.h>

#include "audio.h"
#include "utils.h"
#include "retro.h"

#include "libretro.h"

#define AUDIO_SAMPLE_ALIGN(s) (((s) + 63) & ~63)
#define AUDIO_SAMPLE_TRUNCATE(s) ((s) & ~63)

#define AUDIO_MAX_VOLUME 0x8000
//#define AUDIO_MAX_FRAME_SIZE 16384

#define AUDIO_CHANNELS 1
#define AUDIO_DEFAULT_SAMPLE_RATE 48000u
#define AUDIO_OUTPUT_COUNT 512u
#define AUDIO_BUFFER_COUNT (AUDIO_OUTPUT_COUNT * 10)

typedef struct
{
    int16_t channel;
} VitaMonoSample;

typedef struct
{
    int16_t left;
    int16_t right;
} VitaStereoSample;

typedef struct
{
    int thread_handle;
    int handle;
    int left_volume;
    int right_volume;
    VitaAudioCallback callback;
    void *userdata;
    int16_t *sample_buffer[2];
    unsigned int samples[2];
} VitaAudioChannelInfo;

static int getAudioSampleRate(uint16_t sample_rate);
static int initResampler();
static void finishResampler();

static int audio_ready = 0;
static int audio_pause = 1;
static volatile int audio_stop = 0;

static VitaAudioChannelInfo vita_audio_status[AUDIO_CHANNELS];
static SpeexResamplerState *resampler = NULL;
static int audio_need_resample = 0;
static int audio_resample_quality = 1;
static uint32_t audio_sample_rate = AUDIO_DEFAULT_SAMPLE_RATE;
#if defined(AUDIO_MONO)
static int audio_stereo = SCE_AUDIO_OUT_MODE_MONO;
#else
static int audio_stereo = SCE_AUDIO_OUT_MODE_STEREO;
#endif

static int16_t *audio_buffer = NULL;
static uint32_t audio_shortsps;
static uint32_t audio_output_shorts_count, audio_buffer_shorts_count;
static volatile uint32_t audio_write_pos, audio_read_pos;

size_t retroAudioSampleBatchCallback(const int16_t *data, size_t frames)
{
    if (!audio_ready)
        return frames;

    const int16_t *in_buf;
    int16_t *out_buf;
    int16_t *resampe_buf = NULL;
    uint64_t out_size;

    if (audio_need_resample)
    {
        uint64_t in_size = frames * 2;
        out_size = (uint64_t)(in_size * (double)AUDIO_DEFAULT_SAMPLE_RATE / core_system_av_info.timing.sample_rate);
        resampe_buf = (int16_t *)malloc(out_size * sizeof(int16_t));

        speex_resampler_process_int(resampler, 0, (const spx_int16_t *)data, (spx_uint32_t *)&in_size, (spx_int16_t *)resampe_buf, (spx_uint32_t *)&out_size);

        in_buf = resampe_buf;
    }
    else
    {
        in_buf = data;
        out_size = frames * 2;
    }

    out_buf = audio_buffer;

    int i;
    int write_pos = audio_write_pos;
    for (i = 0; i < out_size; i++, write_pos++)
    {
        if (write_pos >= audio_buffer_shorts_count)
            write_pos = 0;
        *(out_buf + write_pos) = *(in_buf + i);
    }
    audio_write_pos = write_pos;

    if (resampe_buf)
        free(resampe_buf);

    return frames;
}

void vitaAudioResetBuffers()
{
    int i, j;
    for (i = 0; i < AUDIO_CHANNELS; i++)
    {
        for (j = 0; j < 2; j++)
        {
            memset(vita_audio_status[i].sample_buffer[j], 0, audio_output_shorts_count * sizeof(int16_t));
        }
    }
    memset(audio_buffer, 0, audio_buffer_shorts_count * sizeof(int16_t));
    audio_write_pos = 0;
    audio_read_pos = 0;
}

void vitaAudioPause()
{
    audio_pause = 1;
}

void vitaAudioResume()
{
    audio_pause = 0;
}

void vitaAudioSetVolume(int channel, int left, int right)
{
    vita_audio_status[channel].left_volume = left;
    vita_audio_status[channel].right_volume = right;
}

int vitaAudioSetFrequency(int channel, int samples, uint16_t freq, int stereo)
{
    return sceAudioOutSetConfig(vita_audio_status[channel].handle, samples, freq, stereo);
}

int vitaAudioResetFrequency(uint32_t sample_rate)
{
    if (!audio_ready || sample_rate == 0)
        return -1;

    int ret = 0;
    APP_LOG("vitaAudioResetFrequency..\n");
    vitaAudioPause();
    if (resampler)
        finishResampler();
    if (!getAudioSampleRate(sample_rate))
        initResampler();
    APP_LOG("audio_sample_rate: %d\n", audio_sample_rate);
    APP_LOG("audio_need_resample: %d\n", audio_need_resample);
    ret = vitaAudioSetFrequency(0, AUDIO_OUTPUT_COUNT, audio_sample_rate, audio_stereo);
    vitaAudioResetBuffers();
    vitaAudioResume();
    APP_LOG("vitaAudioResetFrequency done\n");

    return ret;
}

void vitaAudioSetChannelCallback(int channel, VitaAudioCallback callback, void *userdata)
{
    volatile VitaAudioChannelInfo *pci = &vita_audio_status[channel];
    pci->callback = callback;
    pci->userdata = userdata;
}

static int vitaAudioOutBlocking(unsigned int channel, unsigned int vol1, unsigned int vol2, void *buf, int length)
{
    if (!audio_ready)
        return -1;

    if (channel >= AUDIO_CHANNELS)
        return -1;

    if (vol1 > AUDIO_MAX_VOLUME)
        vol1 = AUDIO_MAX_VOLUME;

    if (vol2 > AUDIO_MAX_VOLUME)
        vol2 = AUDIO_MAX_VOLUME;

    int vols[2] = {vol1, vol2};
    sceAudioOutSetConfig(vita_audio_status[channel].handle, length, -1, -1);
    sceAudioOutSetVolume(vita_audio_status[channel].handle, SCE_AUDIO_VOLUME_FLAG_L_CH | SCE_AUDIO_VOLUME_FLAG_R_CH, vols);

    return sceAudioOutOutput(vita_audio_status[channel].handle, buf);
}

int vitaAudioCallback(void *buf, unsigned int samples, void *userdata)
{
    int i;
    int16_t *in_buf = (int16_t *)userdata;
    int16_t *out_buf = (int16_t *)buf;

    int loop = audio_read_pos > audio_write_pos;
    for (i = 0; i < samples * audio_shortsps; i++, audio_read_pos++)
    {
        if (audio_read_pos >= audio_buffer_shorts_count)
        {
            audio_read_pos = 0;
            loop = 0;
        }
        if (!loop && audio_read_pos >= audio_write_pos)
        {
            int j;
            for (j = i; j < samples * audio_shortsps; j++)
                *(out_buf + j) = 0;
            return i;
        }
        *(out_buf + i) = *(in_buf + audio_read_pos);
        *(in_buf + audio_read_pos) = 0;
    }

    return samples;
}

static int vitaAudioChannelThread(unsigned int args, void *argp)
{
    APP_LOG("vitaAudioChannelThread start\n");

    volatile int bufidx = 0;
    int channel = *(int *)argp;
    int i;
    int len;
    int16_t *ptr_s;
    void *buffer;
    unsigned int samples;
    VitaAudioCallback callback;

    while (!audio_stop)
    {
        if (audio_pause)
        {
            sceKernelDelayThread(1000);
            continue;
        }

        len = 0;
        callback = vita_audio_status[channel].callback;
        buffer = vita_audio_status[channel].sample_buffer[bufidx];
        samples = vita_audio_status[channel].samples[bufidx];

        if (callback)
            len = callback(buffer, samples, vita_audio_status[channel].userdata);

        if (len <= 0)
            memset(buffer, 0, samples * audio_shortsps * sizeof(int16_t));

        vitaAudioOutBlocking(channel, vita_audio_status[channel].left_volume,
                             vita_audio_status[channel].right_volume, buffer, samples);

        bufidx = (bufidx ? 0 : 1);
    }

    APP_LOG("vitaAudioChannelThread exit\n");
    sceKernelExitThread(0);
    return 0;
}

static int getAudioSampleRate(uint16_t sample_rate)
{
    uint32_t sample_rates[] = {8000, 11025, 12000, 16000, 22050, 24000, 32000, 44100, 48000};
    int i;
    int neglect = 50;
    for (i = 0; i < sizeof(sample_rates) / sizeof(uint32_t); i++)
    {
        if ((sample_rate >= sample_rates[i] - neglect) && (sample_rate <= sample_rates[i] + neglect))
        {
            audio_sample_rate = sample_rates[i];
            return 1;
        }
    }
    
    audio_sample_rate = AUDIO_DEFAULT_SAMPLE_RATE;
    return 0;
}

static int vitaAudioOutputInit()
{
    audio_need_resample = 0;
    audio_write_pos = 0;
    audio_read_pos = 0;

    if (audio_stereo == SCE_AUDIO_OUT_MODE_STEREO)
        audio_shortsps = sizeof(VitaStereoSample) / sizeof(int16_t);
    else
        audio_shortsps = sizeof(VitaMonoSample) / sizeof(int16_t);

    audio_output_shorts_count = AUDIO_OUTPUT_COUNT * audio_shortsps;
    audio_buffer_shorts_count = AUDIO_BUFFER_COUNT * audio_shortsps;

    audio_buffer = (int16_t *)malloc(audio_buffer_shorts_count * sizeof(int16_t));
    if (!audio_buffer)
    {
        vitaAudioShutdown();
        sceKernelExitProcess(0);
        return 0;
    }
    memset(audio_buffer, 0, audio_buffer_shorts_count * sizeof(int16_t));

    APP_LOG("system_sample_rate: %.0f\n", core_system_av_info.timing.sample_rate);
    if (!getAudioSampleRate((uint32_t)core_system_av_info.timing.sample_rate))
        initResampler();
    APP_LOG("audio_need_resample: %d\n", audio_need_resample);

    int i, j;
    for (i = 0; i < AUDIO_CHANNELS; i++)
    {
        vita_audio_status[i].handle = sceAudioOutOpenPort(SCE_AUDIO_OUT_PORT_TYPE_BGM, AUDIO_OUTPUT_COUNT,
                                                          audio_sample_rate, audio_stereo);
        if (vita_audio_status[i].handle < 0)
        {
            vita_audio_status[i].handle = -1;
            vitaAudioShutdown();
            sceKernelExitProcess(0);
            return 0;
        }

        for (j = 0; j < 2; j++)
        {
            vita_audio_status[i].sample_buffer[j] = (int16_t *)malloc(AUDIO_OUTPUT_COUNT * sizeof(int16_t) * audio_shortsps);
            if (!vita_audio_status[i].sample_buffer[j])
            {
                vitaAudioShutdown();
                sceKernelExitProcess(0);
                return 0;
            }
            memset(vita_audio_status[i].sample_buffer[j], 0, AUDIO_OUTPUT_COUNT * sizeof(int16_t) * audio_shortsps);
            vita_audio_status[i].samples[j] = AUDIO_OUTPUT_COUNT;
        }
    }
    vitaAudioSetChannelCallback(0, vitaAudioCallback, audio_buffer);

    return 1;
}

static int vitaAudioThreadInit()
{
    audio_stop = 0;
    audio_pause = 1;

    char label[16];
    strcpy(label, "audiotX");

    int i;
    for (i = 0; i < AUDIO_CHANNELS; i++)
    {
        label[6] = '0' + i;
        vita_audio_status[i].thread_handle = sceKernelCreateThread(label, vitaAudioChannelThread,
                                                                   0x10000100, 0x10000, 0, 0, NULL);

        if (vita_audio_status[i].thread_handle < 0)
        {
            vita_audio_status[i].thread_handle = -1;
            vitaAudioShutdown();
            sceKernelExitProcess(0);
            return 0;
        }

        if (sceKernelStartThread(vita_audio_status[i].thread_handle, sizeof(i), &i) != 0)
        {
            vitaAudioShutdown();
            sceKernelExitProcess(0);
            return 0;
        }
    }

    return 1;
}

int vitaAudioInit()
{
    APP_LOG("vitaAudioInit...\n");

    if (audio_ready)
        vitaAudioShutdown();
    if (resampler)
        finishResampler();

    audio_ready = 0;
    audio_pause = 1;
    audio_stop = 0;

    int i, j;
    for (i = 0; i < AUDIO_CHANNELS; i++)
    {
        vita_audio_status[i].handle = -1;
        vita_audio_status[i].thread_handle = -1;
        vita_audio_status[i].left_volume = AUDIO_MAX_VOLUME;
        vita_audio_status[i].right_volume = AUDIO_MAX_VOLUME;
        vita_audio_status[i].callback = NULL;
        vita_audio_status[i].userdata = NULL;

        for (j = 0; j < 2; j++)
        {
            vita_audio_status[i].sample_buffer[j] = NULL;
            vita_audio_status[i].samples[j] = 0;
        }
    }

    if (!vitaAudioOutputInit())
        return -1;

    if (!vitaAudioThreadInit())
        return -1;

    audio_ready = 1;

    APP_LOG("vitaAudioInit done\n");

    return AUDIO_OUTPUT_COUNT;
}

static void vitaAudioOutputShutdown()
{
    int i, j;
    for (i = 0; i < AUDIO_CHANNELS; i++)
    {
        if (vita_audio_status[i].handle != -1)
        {
            sceAudioOutReleasePort(vita_audio_status[i].handle);
            vita_audio_status[i].handle = -1;
        }

        for (j = 0; j < 2; j++)
        {
            if (vita_audio_status[i].sample_buffer[j])
            {
                free(vita_audio_status[i].sample_buffer[j]);
                vita_audio_status[i].sample_buffer[j] = NULL;
            }
        }
    }

    if (audio_buffer)
    {
        free(audio_buffer);
        audio_buffer = NULL;
    }
}

static void vitaAudioThreadShutdown()
{
    audio_pause = 1;
    audio_stop = 1;

    int i;
    for (i = 0; i < AUDIO_CHANNELS; i++)
    {
        if (vita_audio_status[i].thread_handle != -1)
        {
            sceKernelWaitThreadEnd(vita_audio_status[i].thread_handle, NULL, NULL);
            sceKernelDeleteThread(vita_audio_status[i].thread_handle);
            vita_audio_status[i].thread_handle = -1;
        }
    }
}

int vitaAudioShutdown()
{
    APP_LOG("vitaAudioShutdown...\n");

    audio_ready = 0;
    audio_pause = 1;
    audio_stop = 1;

    vitaAudioThreadShutdown();
    vitaAudioOutputShutdown();
    finishResampler();

    APP_LOG("vitaAudioShutdown done\n");

    return 0;
}

static int initResampler()
{
    if (resampler)
        finishResampler();

    APP_LOG("initResampler...\n");
    resampler = speex_resampler_init(1, (uint32_t)core_system_av_info.timing.sample_rate, audio_sample_rate, audio_resample_quality, NULL);
    if (!resampler)
    {
        vitaAudioShutdown();
        sceKernelExitProcess(0);
        return 0;
    }
    speex_resampler_skip_zeros(resampler);
    audio_need_resample = 1;
    APP_LOG("initResampler done\n");

    return 1;
}

static void finishResampler()
{
    APP_LOG("finishResampler...\n");
    if (resampler)
        speex_resampler_destroy(resampler);
    resampler = NULL;
    audio_need_resample = 0;
    APP_LOG("finishResampler done\n");
}
